from xml.dom.minidom import Document
import cv2
import os
import glob
import shutil
import numpy as np

def generate_xml(name, lines, img_size, class_sets, doncateothers=True):
    doc = Document()

    def append_xml_node_attr(child, parent=None, text=None):
        ele = doc.createElement(child)
        if not text is None:
            text_node = doc.createTextNode(text)
            ele.appendChild(text_node)
        parent = doc if parent is None else parent
        parent.appendChild(ele)
        return ele

    img_name = name + '.jpg'
    # create header
    annotation = append_xml_node_attr('annotation')
    append_xml_node_attr('folder', parent=annotation, text='text')
    append_xml_node_attr('filename', parent=annotation, text=img_name)
    source = append_xml_node_attr('source', parent=annotation)
    append_xml_node_attr('database', parent=source, text='coco_text_database')
    append_xml_node_attr('annotation', parent=source, text='text')
    append_xml_node_attr('image', parent=source, text='text')
    append_xml_node_attr('flickrid', parent=source, text='000000')
    owner = append_xml_node_attr('owner', parent=annotation)
    append_xml_node_attr('name', parent=owner, text='ms')
    size = append_xml_node_attr('size', annotation)
    append_xml_node_attr('width', size, str(img_size[1]))
    append_xml_node_attr('height', size, str(img_size[0]))
    append_xml_node_attr('depth', size, str(img_size[2]))
    append_xml_node_attr('segmented', parent=annotation, text='0')

    # create objects
    objs = []
    for line in lines:
        splitted_line = line.strip().lower().split()
        cls = splitted_line[0].lower()  # 切片后第一个元素：cls: 字符串, "text"

        if not doncateothers and cls not in class_sets:  # ??????
            continue
        # 不关心的类别，继续下一个line
        cls = 'dontcare' if cls not in class_sets else cls
        if cls == 'dontcare':
            continue

        obj = append_xml_node_attr('object', parent=annotation)
        occlusion = int(0)
        x1, y1, x2, y2 = int(float(splitted_line[1]) + 1), int(float(splitted_line[2]) + 1), \
                         int(float(splitted_line[3]) + 1), int(float(splitted_line[4]) + 1)
        truncation = float(0)
        # ？？？difficult=False forever
        difficult = 1 if _is_hard(cls, truncation, occlusion, x1, y1, x2, y2) else 0
        truncted = 0 if truncation < 0.5 else 1

        append_xml_node_attr('name', parent=obj, text=cls)
        append_xml_node_attr('pose', parent=obj, text='none')
        append_xml_node_attr('truncated', parent=obj, text=str(truncted))
        append_xml_node_attr('difficult', parent=obj, text=str(int(difficult)))
        bb = append_xml_node_attr('bndbox', parent=obj)
        append_xml_node_attr('xmin', parent=bb, text=str(x1))
        append_xml_node_attr('ymin', parent=bb, text=str(y1))
        append_xml_node_attr('xmax', parent=bb, text=str(x2))
        append_xml_node_attr('ymax', parent=bb, text=str(y2))

        o = {'class': cls, 'box': np.asarray([x1, y1, x2, y2], dtype=float), \
             'truncation': truncation, 'difficult': difficult, 'occlusion': occlusion}
        objs.append(o)

    return doc, objs


def _is_hard(cls, truncation, occlusion, x1, y1, x2, y2):
    hard = False
    if y2 - y1 < 25 and occlusion >= 2:
        hard = True
        return hard
    if occlusion >= 3:
        hard = True
        return hard
    if truncation > 0.8:
        hard = True
        return hard
    return hard


def build_voc_dirs(outdir):
    mkdir = lambda dir: os.makedirs(dir) if not os.path.exists(dir) else None
    mkdir(outdir)
    mkdir(os.path.join(outdir, 'Annotations'))
    mkdir(os.path.join(outdir, 'ImageSets'))
    mkdir(os.path.join(outdir, 'ImageSets', 'Layout'))
    mkdir(os.path.join(outdir, 'ImageSets', 'Main'))
    mkdir(os.path.join(outdir, 'ImageSets', 'Segmentation'))
    mkdir(os.path.join(outdir, 'JPEGImages'))
    mkdir(os.path.join(outdir, 'SegmentationClass'))
    mkdir(os.path.join(outdir, 'SegmentationObject'))
    return os.path.join(outdir, 'Annotations'), os.path.join(outdir, 'JPEGImages'), os.path.join(outdir, 'ImageSets',
                                                                                                 'Main')


if __name__ == '__main__':
    _outdir = 'TEXTVOC/VOC2007'
    _draw = bool(0)  # False
    _dest_label_dir, _dest_img_dir, _dest_set_dir = build_voc_dirs(_outdir)
    _doncateothers = bool(1)  # True
    for dset in ['train']:
        _labeldir = 'label_tmp'
        _imagedir = 're_image'
        class_sets = ('text', 'dontcare')
        class_sets_dict = dict((k, i) for i, k in enumerate(class_sets))  # {"text": 0, "dontcare": 1}
        allclasses = {}  # 记录每类box的个数，键值为cls和数量num，一张图片有多个box
        # "ImageSets/Main/text_train.txt", "ImageSets/Main/dontcare_train.txt", 以 "w" 打开这两个文件，在list中
        fs = [open(os.path.join(_dest_set_dir, cls + '_' + dset + '.txt'), 'w') for cls in class_sets]
        # "ImageSets/Main/train.txt", "w"方式打开，存放图片路径
        ftrain = open(os.path.join(_dest_set_dir, dset + '.txt'), 'w')

        # 列出 "./label_tmp/*.txt" 文件，返回一个list, 该txt文件包含多个小框的四个坐标
        files = glob.glob(os.path.join(_labeldir, '*.txt'))
        files.sort()
        for file in files:  # 每个file也就是包含boxes信息的txt文件
            # 分离：路径/文件名
            path, basename = os.path.split(file)
            # 分离：name/扩展名
            stem, ext = os.path.splitext(basename)
            with open(file, 'r') as f:
                lines = f.readlines()  # 每个line为一个anchor的四个坐标信息
            # "re_image/" + name + ".jpg", 就是resize后的图片
            img_file = os.path.join(_imagedir, stem + '.jpg')

            print(img_file)
            img = cv2.imread(img_file)
            img_size = img.shape

            # doc：xml文件；objs：包含多个字典的list，每个字典就是一张图片里每个box的信息
            doc, objs = generate_xml(stem, lines, img_size, class_sets=class_sets, doncateothers=_doncateothers)

            # 保存resize的训练图片和label的xml文件
            cv2.imwrite(os.path.join(_dest_img_dir, stem + '.jpg'), img)
            xmlfile = os.path.join(_dest_label_dir, stem + '.xml')
            with open(xmlfile, 'w') as f:
                f.write(doc.toprettyxml(indent='	'))

            # 位于"./ImageSets/main"下的"train.txt"文件，写入图片名
            ftrain.writelines(stem + '\n')  # 多行写入f.writelines(["str1 \n", "str2 \n", ...])

            cls_in_image = set([o['class'] for o in objs]) # maybe ("text")

            for obj in objs:  # 遍历每个字典，每个字典就是一张图片里每个box的信息
                cls = obj['class']
                # 该cls不在字典中，key：cls，value：0；若在：value + 1
                allclasses[cls] = 0 \
                    if not cls in list(allclasses.keys()) else allclasses[cls] + 1

            # 在"./ImageSets/Main/text_train.txt"写入imageName 1 or -1
            # 在每一类的txt文件中记录img的name和是否包含此类的信息
            for cls in cls_in_image:
                if cls in class_sets:
                    fs[class_sets_dict[cls]].writelines(stem + ' 1\n')
            for cls in class_sets:
                if cls not in cls_in_image:
                    fs[class_sets_dict[cls]].writelines(stem + ' -1\n')


        (f.close() for f in fs)
        ftrain.close()

        print('~~~~~~~~~~~~~~~~~~~')
        print(allclasses)
        print('~~~~~~~~~~~~~~~~~~~')
        # 将train.txt中的图像名copy到val.txt和trainval.txt中
        shutil.copyfile(os.path.join(_dest_set_dir, 'train.txt'), os.path.join(_dest_set_dir, 'val.txt'))
        shutil.copyfile(os.path.join(_dest_set_dir, 'train.txt'), os.path.join(_dest_set_dir, 'trainval.txt'))
        for cls in class_sets:
            shutil.copyfile(os.path.join(_dest_set_dir, cls + '_train.txt'),
                            os.path.join(_dest_set_dir, cls + '_trainval.txt'))
            shutil.copyfile(os.path.join(_dest_set_dir, cls + '_train.txt'),
                            os.path.join(_dest_set_dir, cls + '_val.txt'))
